# Copyright (c) 2024-present AI-Labs

from fastapi import APIRouter

import base64
import os
import torch
import uuid

from datetime import datetime

from melotts.api import TTS

"""
初始化模型
"""
model = TTS(language="ZH", device="auto")
speaker_ids = model.hps.data.spk2id

"""
路由信息设置
"""
router = APIRouter(
    prefix='/tts/melotts',
    tags = ['语音生成']
)

"""
文本数据转语音
"""
def tts_to_file(input):
    # 创建一个本地目录
    localdir = f"audio/{datetime.now().strftime('%Y-%m-%d')}"
    os.makedirs(f"{config.setting.statics.path}/{localdir}", exist_ok=True)
    # 指定要保存到本地的文件名
    localfile = f"{localdir}/{uuid.uuid4()}.wav"
    
    # 将文本转换成音频文件，并保存到本地文件
    model.tts_to_file(input, speaker_ids['ZH'], f"{config.setting.statics.path}/{localfile}", speed=0.7)

    # 清理显存缓存
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()
    
    return localfile


"""
文本数据转语音
"""
@router.post("/v1")
def generate_audio_post(request: dict):
    # 将文本数据转换成语音，保存到本地文件
    localfile = tts_to_file(request["text"])

    # 获取本地文件的字节数据信息
    with open(f"{config.setting.statics.path}/{localfile}", 'rb') as audio_file:
        """
        读取文件内容到字节字符串
        """
        audio_bytes = audio_file.read()
    
    # 将字节数据做Base64编码，编码成字符串，并返回
    return {"tts_result": base64.b64encode(audio_bytes).decode('utf-8')}
